import tensorflow as tf
import tensorflow.keras as keras
from tensorflow.keras import layers, activations, losses, metrics, optimizers, callbacks
import numpy as np
import os


def ConvBnRelu(ch, ksize, strides=1):
    return keras.Sequential([
        layers.Conv2D(ch, ksize, strides, 'same'),
        layers.BatchNormalization(),
        layers.ReLU()
    ])


class InceptionBlock(keras.Model):

    def __init__(self, branch_ch, shrink=False, **kwargs):
        super().__init__(**kwargs)
        if shrink:
            strides = 2
        else:
            strides = 1
        self.branch01 = ConvBnRelu(branch_ch, 1, strides)
        self.branch02_1 = ConvBnRelu(branch_ch, 1, strides)
        self.branch02_2 = ConvBnRelu(branch_ch, 3)
        self.branch03_1 = ConvBnRelu(branch_ch, 1, strides)
        self.branch03_2 = ConvBnRelu(branch_ch, 5)
        self.branch04_1 = layers.MaxPool2D((3, 3), strides=(1, 1), padding='same')
        self.branch04_2 = ConvBnRelu(branch_ch, 1, strides)

    def call(self, inputs, training=None, mask=None):
        x1 = self.branch01(inputs, training=training)
        x2 = self.branch02_1(inputs, training=training)
        x2 = self.branch02_2(x2, training=training)
        x3 = self.branch03_1(inputs, training=training)
        x3 = self.branch03_2(x3, training=training)
        x4 = self.branch04_1(inputs, training=training)
        x4 = self.branch04_2(x4, training=training)
        x = tf.concat([x1, x2, x3, x4], axis=3)
        return x


class InceptionNet(keras.Model):

    def __init__(self, n_cls, init_out, n_blocks, act='softmax', **kwargs):
        super().__init__(**kwargs)
        self.conv = ConvBnRelu(init_out, 3)

        self.blocks = keras.Sequential()
        ch = init_out
        for block_id in range(n_blocks):
            for layer_id in range(2):
                if layer_id == 0:
                    shrink = True
                    ch *= 2
                else:
                    shrink = False
                block = InceptionBlock(ch, shrink)
                self.blocks.add(block)

        self.glb_avg = layers.GlobalAvgPool2D()
        self.fc = layers.Dense(n_cls, activation=act)

    def call(self, inputs, training=None, mask=None):
        x = self.conv(inputs, training=training)
        x = self.blocks(x)
        x = self.glb_avg(x)
        x = self.fc(x)
        return x


if '__main__' == __name__:
    BATCH_SIZE = 64
    N_EPOCHS = 3
    FILE_NAME = os.path.basename(__file__)
    LOG_DIR = os.path.join('_log', FILE_NAME)

    model = InceptionNet(10, 16, 2, act=None)
    model.build(input_shape=(None, 32, 32, 3))
    model.summary()

    np.random.seed(777)
    tf.random.set_seed(777)

    (x_train, y_train), (x_test, y_test) = keras.datasets.cifar10.load_data()

    x_train = x_train.astype(np.float32) / 255.
    x_test = x_test.astype(np.float32) / 255.

    model.compile(
        loss=losses.SparseCategoricalCrossentropy(from_logits=True),
        optimizer=optimizers.Adam(),
        metrics=metrics.sparse_categorical_accuracy,
    )

    res = model.fit(x_train, y_train,
                    batch_size=BATCH_SIZE, epochs=N_EPOCHS,
                    validation_batch_size=BATCH_SIZE, validation_split=0.1,
                    callbacks=callbacks.TensorBoard(log_dir=LOG_DIR, update_freq='batch', profile_batch=0))
    print(res)

    model.evaluate(x_test, y_test, batch_size=BATCH_SIZE)
